import psycopg2
import os

DB_NAME = 'game_trading_v2'         # 你要创建的新数据库名
DB_USER = 'postgres'
DB_PASSWORD = '145114'
DB_HOST = 'localhost'
DB_PORT = '5432'
DATA_DIR = '../data'

def create_database():
    # 连接到默认数据库
    conn = psycopg2.connect(
        dbname='postgres',
        user=DB_USER,
        password=DB_PASSWORD,
        host=DB_HOST,
        port=DB_PORT
    )
    conn.autocommit = True
    cur = conn.cursor()
    # 检查数据库是否存在，不存在则创建
    cur.execute(f"SELECT 1 FROM pg_database WHERE datname='{DB_NAME}'")
    exists = cur.fetchone()
    if not exists:
        cur.execute(f'CREATE DATABASE "{DB_NAME}"')
        print(f"数据库 {DB_NAME} 已创建。")
    else:
        print(f"数据库 {DB_NAME} 已存在。")
    cur.close()
    conn.close()

def import_csv():
    create_table_sql = [
        # 用户表，UserID自增
        """
        CREATE TABLE IF NOT EXISTS Users (
            UserID SERIAL PRIMARY KEY,
            Username VARCHAR(50) UNIQUE NOT NULL,
            Password VARCHAR(64) NOT NULL,
            Email VARCHAR(100) UNIQUE NOT NULL,
            Phone VARCHAR(20),
            RegistrationDate TIMESTAMP,
            Balance NUMERIC(10,2) DEFAULT 0,
            CreditRating INT,
            Status VARCHAR(20),
            Role VARCHAR(20) DEFAULT 'user'
        );
        """,
        # 装备表，EquipmentID自增
        """
        CREATE TABLE IF NOT EXISTS Equipments (
            EquipmentID SERIAL PRIMARY KEY,
            SellerID INT REFERENCES Users(UserID),
            EquipmentName VARCHAR(100) NOT NULL,
            Category VARCHAR(20),
            Level INT,
            Rarity VARCHAR(20),
            Attributes JSON,
            Price NUMERIC(10,2) NOT NULL,
            ListedDate TIMESTAMP,
            Status VARCHAR(20),
            ImageURL VARCHAR(255)
        );
        """,
        # 交易表，TransactionID自增
        """
        CREATE TABLE IF NOT EXISTS Transactions (
            TransactionID SERIAL PRIMARY KEY,
            EquipmentID INT REFERENCES Equipments(EquipmentID),
            SellerID INT REFERENCES Users(UserID),
            BuyerID INT REFERENCES Users(UserID),
            TransactionDate TIMESTAMP,
            Price NUMERIC(10,2),
            Status VARCHAR(30),
            PaymentMethod VARCHAR(20),
            CompletionDate TIMESTAMP
        );
        """,
        # 钱包流水表，RecordID自增
        """
        CREATE TABLE IF NOT EXISTS WalletRecords (
            RecordID SERIAL PRIMARY KEY,
            UserID INT REFERENCES Users(UserID),
            TransactionID INT REFERENCES Transactions(TransactionID),
            Type VARCHAR(20),
            Amount NUMERIC(10,2),
            Balance NUMERIC(10,2),
            Timestamp TIMESTAMP
        );
        """,
        # 购物车表，CartID自增
        """
        CREATE TABLE IF NOT EXISTS Cart (
            CartID SERIAL PRIMARY KEY,
            UserID INT REFERENCES Users(UserID),
            EquipmentID INT REFERENCES Equipments(EquipmentID),
            AddedDate TIMESTAMP
        );
        """,
        # 价格变动消息表，MessageID自增
        """
        CREATE TABLE IF NOT EXISTS PriceChangeMessages (
            MessageID SERIAL PRIMARY KEY,
            UserID INT REFERENCES Users(UserID),
            EquipmentID INT REFERENCES Equipments(EquipmentID),
            OldPrice NUMERIC(10,2),
            NewPrice NUMERIC(10,2),
            ChangeDate TIMESTAMP,
            IsRead BOOLEAN
        );
        """
    ]
    table_files = [
        ('Users', 'Users.csv'),
        ('Equipments', 'Equipments.csv'),
        ('Transactions', 'Transactions.csv'),
        ('WalletRecords', 'WalletRecords.csv'),
        ('Cart', 'Cart.csv'),
        ('PriceChangeMessages', 'PriceChangeMessages.csv')
    ]
    conn = psycopg2.connect(
        dbname=DB_NAME,
        user=DB_USER,
        password=DB_PASSWORD,
        host=DB_HOST,
        port=DB_PORT
    )
    cur = conn.cursor()
    # 先清空所有表，避免外键冲突（注意顺序）
    for table in ['PriceChangeMessages', 'Cart', 'WalletRecords', 'Transactions', 'Equipments', 'Users']:
        cur.execute(f"DROP TABLE IF EXISTS {table} CASCADE;")
    conn.commit()
    # 创建表
    for sql in create_table_sql:
        cur.execute(sql)
    conn.commit()

    # 设置所有自增ID从100000开始
    seqs = [
        'users_userid_seq',
        'equipments_equipmentid_seq',
        'transactions_transactionid_seq',
        'walletrecords_recordid_seq',
        'cart_cartid_seq',
        'pricechangemessages_messageid_seq'
    ]
    for seq in seqs:
        cur.execute(f"ALTER SEQUENCE {seq} RESTART WITH 100000;")
    conn.commit()

    # 导入数据
    for table, csv_file in table_files:
        csv_path = os.path.join(DATA_DIR, csv_file)
        if not os.path.exists(csv_path):
            print(f"{csv_path} 不存在，跳过导入 {table}")
            continue
        with open(csv_path, 'r', encoding='utf-8-sig') as f:
            cur.copy_expert(
                f"COPY {table} FROM STDIN WITH CSV HEADER DELIMITER ',' ENCODING 'UTF8'",
                f
            )
        conn.commit()
        print(f"{table} 导入完成")
    cur.close()
    conn.close()
    print("所有数据导入完成。")

if __name__ == '__main__':
    create_database()
    import_csv()